package com.jsn.filter;

import com.alibaba.fastjson.JSONObject;

import com.jsn.responseEntity.CommonCode;
import com.jsn.responseEntity.R;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Map;

//@Component
public class MyRequestRateLimiter extends RequestRateLimiterGatewayFilterFactory {

	private final RateLimiter defaultRateLimiter;

	private final KeyResolver defaultKeyResolver;

	@Autowired
	public MyRequestRateLimiter(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver) {
		super(defaultRateLimiter, defaultKeyResolver);
		this.defaultRateLimiter = defaultRateLimiter;
		this.defaultKeyResolver = defaultKeyResolver;
	}


	@Override
	public GatewayFilter apply(Config config) {


		KeyResolver resolver = getOrDefault(config.getKeyResolver(), defaultKeyResolver);
		RateLimiter<Object> limiter = getOrDefault(config.getRateLimiter(), defaultRateLimiter);

		boolean denyEmpty = (Boolean)this.getOrDefault(config.getDenyEmptyKey(), super.isDenyEmptyKey());
		HttpStatusHolder emptyKeyStatus = HttpStatusHolder.parse((String)this.getOrDefault(config.getEmptyKeyStatus(), super.getEmptyKeyStatusCode()));
		return (exchange, chain) -> {
			return resolver.resolve(exchange).defaultIfEmpty("____EMPTY_KEY__").flatMap((key) -> {
				if ("____EMPTY_KEY__".equals(key)) {
					if (denyEmpty) {
						ServerWebExchangeUtils.setResponseStatus(exchange, emptyKeyStatus);
						return exchange.getResponse().setComplete();
					} else {
						return chain.filter(exchange);
					}
				} else {
					String routeId = config.getRouteId();
					if (routeId == null) {
						Route route = (Route)exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
						routeId = route.getId();
					}

					return limiter.isAllowed(routeId, key).flatMap((response) -> {
						Iterator var4 = response.getHeaders().entrySet().iterator();

						while(var4.hasNext()) {
							Map.Entry<String, String> header = (Map.Entry)var4.next();
							exchange.getResponse().getHeaders().add((String)header.getKey(), (String)header.getValue());
						}

						if (response.isAllowed()) {
							return chain.filter(exchange);
						} else {
				/*			ServerWebExchangeUtils.setResponseStatus(exchange, config.getStatusCode());
							return exchange.getResponse().setComplete();*/

							ServerHttpResponse httpResponse = exchange.getResponse();
							httpResponse.setStatusCode(config.getStatusCode());
							if (!httpResponse.getHeaders().containsKey("Content-Type")) {
								httpResponse.getHeaders().add("Content-Type", "application/json");
							}


							DataBuffer buffer = null;

								buffer = httpResponse.bufferFactory().wrap(JSONObject.toJSONString(R.error(CommonCode.ACCESS_LIMIT)).getBytes(StandardCharsets.UTF_8));

							return httpResponse.writeWith(Mono.just(buffer));

						}
					});
				}
			});
		};

	}
	private <T> T getOrDefault(T configValue, T defaultValue) {
		return (configValue != null) ? configValue : defaultValue;
	}
}
